import pandas as pd
import os
from sklearn.model_selection import train_test_split
from sklearn.model_selection import cross_val_score
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB

# 性别预测(有待完善)
# 尝试使用用户昵称来构造分类模型，预测昵称的用户性别。
#
# 数据整理
# 首先导入需要的库，接着合并数据，然后筛选出用户性别为男和女的用户。

list_all = os.listdir('infoFriend')
all_list = []
for i in list_all:
    path = 'infoFriend/' + i
    df = pd.read_csv(open(path, encoding='utf_8_sig'))
    all_list.append(df)
all_data = pd.concat([all_list[0]])

df = all_data[(all_data['Sex'] == 1) | (all_data['Sex'] == 2)]

# 数据预处理
# 这里划分数据集，并通过CountVectorizer将数据转换为词向量。
X_train, X_test, Y_train, Y_test = train_test_split(df['NickName'].values.astype('U'), df['Sex'].values.astype('U'),
                                                    test_size=0.2, random_state=22)
count_vect = CountVectorizer()
X_train_cov = count_vect.fit_transform(X_train)

clf = MultinomialNB(alpha=0.0001)
clf.fit(X_train_cov, Y_train)

print(clf.score(X_train_cov, Y_train))

X_test_cov = count_vect.transform(X_test)
print(clf.score(X_test_cov, Y_test))

test = ['。。。', 'abc', 'ace', 'avril']
X = count_vect.transform(test)
print(clf.predict(X))
